
from models import PatchCore
from torch.utils.data import DataLoader
from datasets import mvtec



def patchcore_run(cfg,phase='train',weights=''):

    print('using patchcore')


    category = cfg['normal_class']
    train_batch_size = cfg['train_batch_size']
    test_batch_size = cfg['test_batch_size']
    load_size = cfg['load_size']
    input_size = cfg['input_size']

    if category == 'all':
        train_class = mvtec.CLASS_NAMES
    else:
        train_class = [category]

    model = PatchCore(cfg)
    # print(train)
    for c in train_class:
        print(c)
        train_dataset = mvtec.MVTecDataset(root_path=cfg['dataset_dir'], class_name=c, is_train=True,resize=load_size, cropsize=input_size)
        train_dataloader = DataLoader(train_dataset, batch_size=train_batch_size, pin_memory=True)
        test_dataset = mvtec.MVTecDataset(root_path=cfg['dataset_dir'], class_name=c, is_train=False,resize=load_size, cropsize=input_size)
        test_dataloader = DataLoader(test_dataset, batch_size=test_batch_size, pin_memory=True)

        if phase == 'train':
            model.train(train_dataloader)
            model.train_after(c)
            model.init_results_list()
            
            model.test_after(test_dataloader,c)
            model.evaluate(c)
            model.init_results_list()
        elif phase == 'test':
            model.test_after(test_dataloader,c)
            model.evaluate(c)
            model.init_results_list()

